
import pybullet
import pybullet_data
import time
import numpy as np
import matplotlib.pyplot as plt



physicsClient = pybullet.connect(pybullet.GUI) #or p.DIRECT for non-graphical version
pybullet.setAdditionalSearchPath(pybullet_data.getDataPath()) # build in data
pybullet.resetDebugVisualizerCamera(cameraDistance=3, cameraYaw=90,
                                    cameraPitch=-5, cameraTargetPosition=[0, 2, 1.])
pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING,0)
pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI,0)
# pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_TINY_RENDERER,0)
pybullet.setGravity(0,0,-9.81)


planeId = pybullet.loadURDF("plane.urdf")
urdfFlags = pybullet.URDF_USE_IMPLICIT_CYLINDER 
#urdfFlags = pybullet.URDF_USE_SELF_COLLISION_EXCLUDE_PARENT
robotId = pybullet.loadURDF("/home/he/hellobot/PybulletSimulation/Simulation2/urdf/Simulation2.urdf", [0., 0., 0.1], 
                                np.array(pybullet.getQuaternionFromEuler([0., 0., 0.])), 
                            useFixedBase=False, flags=urdfFlags, globalScaling=1)

pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING,1)
jointsNum = pybullet.getNumJoints(robotId)
jointsInfo = []
for i in range(jointsNum):
    jInfo = pybullet.getJointInfo(robotId,i)
    print(jInfo)
    jointsInfo.append(jInfo)

pybullet.changeVisualShape(robotId,2,rgbaColor=[0.1,0.1,0.1,0.1])# This Line Change the color and Transparency of body
sim_frcy = 60
dt = 1./sim_frcy
pybullet.setTimeStep(dt) # default time step 240Hz
sim_t = 0.

Kspring = 1700 #弹簧刚度
Dspring = 0 #弹簧阻尼

springq0 = pybullet.getJointState(robotId, jointsInfo[3][0])[0] 
#pybullet.setJointMotorControl2(robotId, jointsInfo[3][0], pybullet.POSITION_CONTROL, targetPosition=springq0-0.005,force=0)
#springq0 = pybullet.getJointState(robotId, jointsInfo[3][0])[0] 
print(springq0)
pybullet.setRealTimeSimulation(0)



while True:
    sim_t += dt
    # if (sim_t/dt)%100 == 0: 

    if sim_t >2 and sim_t <3:
        q,dq,_,_ = pybullet.getJointState(robotId, jointsInfo[3][0])
       
        springForce = -Kspring * (q- (springq0  - 0.08) ) - Dspring * dq

        pybullet.setJointMotorControl2(robotId, jointsInfo[3][0], pybullet.TORQUE_CONTROL, force=springForce)
        #print(pybullet.getJointState(robotId, jointsInfo[3][0])[0])
    if sim_t > 3: 
        pybullet.setJointMotorControl2(robotId, jointsInfo[3][0], pybullet.POSITION_CONTROL, targetPosition=springq0,force=100)
    print('time%1.4fs'%sim_t)
    print(pybullet.getJointState(robotId, jointsInfo[3][0])[0])

    pybullet.stepSimulation()
    time.sleep(1./sim_frcy)
pybullet.disconnect()

